from uvr5_2.UVR import *
from uvr5_2.gui_data.constants import MDX_ARCH_TYPE

params = {
    "is_primary_stem_only": False,
    "is_secondary_stem_only": False,
    "is_primary_model_primary_stem_only": False,
    "is_primary_model_secondary_stem_only": False,
    "is_pre_proc_model": False,
    "is_secondary_model": False,
    "primary_model_primary_stem": False,
    "denoise_option": None,
    "vocal_opt": "",
    "is_vocal_split_model": False,
}
# 分离人声
def split_vocal(input_path, output_path):

    params["is_primary_stem_only"] = True
    params["is_pre_proc_model"] = False
    params["primary_model_primary_stem"] = False
    params["is_primary_model_primary_stem_only"] = False
    params["is_vocal_split_model"] = False
    params["vocal_opt"] = "Vocals"
    return process_start(
        input_path,
        output_path,
        MDX_ARCH_TYPE,
        "MDX23C-InstVoc HQ",
        params
    )


# 分离伴奏
def split_instrumental(input_path, output_path):
    params["is_secondary_stem_only"] = True
    params["is_pre_proc_model"] = False
    params["primary_model_primary_stem"] = False
    params["is_primary_model_primary_stem_only"] = False
    params["is_vocal_split_model"] = False
    params["vocal_opt"] = "Instrumental"
    return process_start(
        input_path,
        output_path,
        MDX_ARCH_TYPE,
        "MDX23C-InstVoc HQ",
        params
    )

# 降噪
def denoise_no_normal(input_path, output_path):
    params["is_primary_stem_only"] = True
    params["is_pre_proc_model"] = False
    params["primary_model_primary_stem"] = False
    params["is_primary_model_primary_stem_only"] = False
    params["is_vocal_split_model"] = False
    params["vocal_opt"] = "No Echo"
    return process_start(
        input_path,
        output_path,
        VR_ARCH_TYPE,
        "UVR-De-Echo-Normal",
        params
    )

def denoise_normal(input_path, output_path):
    params["is_secondary_stem_only"] = True
    params["is_pre_proc_model"] = False
    params["primary_model_primary_stem"] = False
    params["is_primary_model_primary_stem_only"] = False
    params["is_vocal_split_model"] = False
    params["vocal_opt"] = "Echo"
    return process_start(
        input_path,
        output_path,
        VR_ARCH_TYPE,
        "UVR-De-Echo-Normal",
        params
    )

def denoise(input_path, output_path):
    params["is_primary_stem_only"] = True
    params["is_pre_proc_model"] = False
    params["primary_model_primary_stem"] = False
    params["is_primary_model_primary_stem_only"] = False
    params["is_vocal_split_model"] = False
    params["vocal_opt"] = "Noise"
    return process_start(
        input_path,
        output_path,
        VR_ARCH_TYPE,
        "UVR-DeNoise",
        params
    )

def denoise_no(input_path, output_path):
    params["is_secondary_stem_only"] = True
    params["is_pre_proc_model"] = False
    params["primary_model_primary_stem"] = False
    params["is_primary_model_primary_stem_only"] = False
    params["is_vocal_split_model"] = False
    params["vocal_opt"] = "No Noise"
    return process_start(
        input_path,
        output_path,
        VR_ARCH_TYPE,
        "UVR-DeNoise",
        params
    )


if __name__ == '__main__':
    output_path = '../data/lanying'
    vocal = split_vocal('/Users/zzh/PycharmProjects/AI/data/lanying/huanzou.WAV', output_path)
    normal = denoise_no_normal(vocal['Vocals'], output_path)
    denoise_no = denoise_no(normal['No Echo'], output_path)
    print(denoise_no)
